from major.models_manager import chat_model
from langgraph.prebuilt import create_react_agent
from typing import Annotated # typing是python的基础库，管理各种数据类型
import datetime
from langchain.tools import tool

def current_time():
    """返回时间的工具"""
    print("正在获取时间...")
    return datetime.datetime.now()

@tool(description="加法运算",args_schema={ # 定义方法所需的参数, 若方法参数为空，则设置为{}或省略parameters
                "type": "object",
                "properties": {
                    "a": {"type": "integer", "description": "第一个数字"}, # LLM会根据参数描述在prompt中提取参数
                    "b": {"type": "integer", "description": "第二个数字"},
                },
                "required": ["a", "b"], # 告知LLM这些参数是必须的
            })
def add_numbers(a,b):
    print("正在执行加法运算...")
    return a + b

def multiply_numbers(
    a : Annotated[float, "第一个数"],
    b : Annotated[float, "第二个数"]):
    """乘法运算"""
    print("正在执行乘法运算...")
    return a * b




model = chat_model.get_model()
# 创建Agent，包含两个工具：calculator 和 current_time
agent = create_react_agent(model,
                           tools=[current_time,add_numbers,multiply_numbers],
                           prompt='''你是一个助手''')


response  = agent.invoke({'messages': [{'role': 'user', 'content': '2*6'}]})
print(response['messages'][-1].content)